package com.test.utils.crypto;

import java.nio.charset.StandardCharsets;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.Base64;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import org.slf4j.LoggerFactory;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.test.util.ConfigManager;
import com.test.utils.ConfigManager;

public class AESEncryptionUtil {
    private static final Logger logger = LoggerFactory.getLogger(AESEncryptionUtil.class);

    private static final String ALGORITHM = "AES";
    private static final String TRANSFORMATION = "AES/GCM/NoPadding";
    private static final int KEY_LENGTH = 256;
    private static final int IV_LENGTH = 12;
    private static final int TAG_LENGTH = 16;

    private static final String PRIVATE_KEY_ENV_VAR = "AES_PRIVATE_KEY";

    private static SecretKey secretKey;
    private static final SecureRandom secureRandom = new SecureRandom();

    static {
        initializeKey();
    }

    private static void initializeKey() {
        try {
            ConfigManager config = ConfigManager.getInstance();
            String keyFromEnv = config.getProperty("PRIVATE_KEY");

            if (keyFromEnv != null && !keyFromEnv.trim().isEmpty()) {
                byte[] keyBytes = Base64.getDecoder().decode(keyFromEnv.trim());
                secretKey = new SecretKeySpec(keyBytes, ALGORITHM);
                logger.info("AES key loaded from environment variable: {}", PRIVATE_KEY_ENV_VAR);
            } else {
                secretKey = generateNewKey();
                String encodedKey = Base64.getEncoder().encodeToString(secretKey.getEncoded());
                logger.warn("No AES key found in environment variable '{}'. Generated new key for development use.", PRIVATE_KEY_ENV_VAR);
                logger.warn("Generated AES Key (Base64): {}", encodedKey);
                logger.warn("Set this key in environment variable for production: export {}={}", PRIVATE_KEY_ENV_VAR, encodedKey);
            }
        } catch (Exception e) {
            logger.error("Failed to initialize AES key", e);
            throw new RuntimeException("Failed to initialize AES encryption", e);
        }
    }

    public static SecretKey generateNewKey() {
        try {
            KeyGenerator keyGenerator = KeyGenerator.getInstance(ALGORITHM);
            keyGenerator.init(KEY_LENGTH);
            return keyGenerator.generateKey();
        } catch (NoSuchAlgorithmException e) {
            logger.error("Failed to generate AES key", e);
            throw new RuntimeException("Failed to generate AES key", e);

        }
    }

    public static String getKeyAsBase64() {
        if (secretKey == null) {
            throw new IllegalStateException("Secret key not initialize");
        }
        return Base64.getEncoder().encodeToString(secretKey.getEncoded());
    }

    public static String encrypt(String plaintext) {
        if (plaintext == null || plaintext.isEmpty()) {
            throw new IllegalArgumentException("Plain text cannot be null or empty");
        }

        try {
            byte[] iv = new byte[IV_LENGTH];
            secureRandom.nextBytes(iv);

            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
            GCMParameterSpec gcmSpec = new GCMParameterSpec(TAG_LENGTH * 8, iv);
            cipher.init(Cipher.ENCRYPT_MODE, secretKey, gcmSpec);

            byte[] encryptedBytes = cipher.doFinal(plaintext.getBytes(StandardCharsets.UTF_8));

            byte[] encryptedWIthIv = new byte[IV_LENGTH + encryptedBytes.length];
            System.arraycopy(iv, 0, encryptedWIthIv, 0, IV_LENGTH);
            System.arraycopy(encryptedBytes, 0, encryptedWIthIv, IV_LENGTH, encryptedBytes.length);

            String result = Base64.getEncoder().encodeToString(encryptedWIthIv);
            logger.debug("Successfully encrypted text of length: {}", plaintext.length());
            return result;
        } catch (Exception e) {
            logger.error("Failed to encrypt text", e);
            throw new RuntimeException("Encryption failed", e);
        }
    }

    public static String decrypt(String encryptedText) {
        if (encryptedText == null || encryptedText.isEmpty()) {
            throw new IllegalArgumentException("Encrypted text cannot be null or empty");
        }

        try {
            byte[] encryptedWithIv = Base64.getDecoder().decode(encryptedText);

            if (encryptedWithIv.length < IV_LENGTH + TAG_LENGTH) {
                throw new IllegalArgumentException("Invalid encrypted text format");
            }

            byte[] iv = new byte[IV_LENGTH];
            byte[] encryptedBytes = new byte[encryptedWithIv.length - IV_LENGTH];
            System.arraycopy(encryptedWithIv, 0, iv, 0, IV_LENGTH);
            System.arraycopy(encryptedWithIv, IV_LENGTH, encryptedBytes, 0, encryptedBytes.length);

            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
            GCMParameterSpec gcmSpec = new GCMParameterSpec(TAG_LENGTH * 8, iv);
            cipher.init(Cipher.DECRYPT_MODE, secretKey, gcmSpec);

            byte[] decryptedBytes = cipher.doFinal(encryptedBytes);
            String result = new String(decryptedBytes, StandardCharsets.UTF_8);

            logger.debug("Successfully decrypted text of length: {}", result.length());
            return result;
        } catch (Exception e) {
            logger.error("Failed to decrypt text", e);
            throw new RuntimeException("Decryption failed", e);
        }
    }

    public static String encryptPassword(String password) {
        return encrypt(password);
    }

    public static String decryptPassword(String encryptedPassword) {
        return decrypt(encryptedPassword);
    }

    public static boolean isInitialized() {
        return secretKey != null;
    }

    public static void reinitialize() {
        secretKey = null;
        initializeKey();
    }

    public static void setKeyFromBase64(String base64Key) {
        if (base64Key == null || base64Key.trim().isEmpty()) {
            throw new IllegalArgumentException("Key cannot be null or empty");
        }

        try {
            byte[] keyBytes = Base64.getDecoder().decode(base64Key.trim());
            secretKey = new SecretKeySpec(keyBytes, ALGORITHM);
            logger.info("Custom AES key set successfully");
        } catch (Exception e) {
            logger.error("Failed to set custom key", e);
            throw new RuntimeException("Failed to set custom key", e);
        }
    }
}
